import os
import docx
from openpyxl import Workbook
#############################################################
# 遍历当前目录下所有docx，把docx中表格的第一列为HEADERS中包含内容的行
# 输出到excel中，并把对应行内容的第二列按照对应关系输出到表格中
# docx中每个表格作为一列
#############################################################
# 定义表头
HEADERS = ["漏洞名称", "影响主机", "漏洞类型", "CVSS评分", "漏洞描述", "修复建议"]

def extract_tables_from_word(word_file):
    """
    从 Word 文档中提取表格数据
    """
    doc = docx.Document(word_file)
    tables = []
    
    # 遍历文档中的所有表格
    for table in doc.tables:
        table_data = {}
        # 遍历表格中的每一行
        for row in table.rows:
            if len(row.cells) == 2:  # 确保表格只有两列
                key = row.cells[0].text.strip()  # 第一列作为键
                value = row.cells[1].text.strip()  # 第二列作为值
                if key in HEADERS:  # 只处理表头中定义的键
                    table_data[key] = value
        tables.append(table_data)
    
    return tables

def write_tables_to_excel(tables, ws):
    """
    将表格数据写入 Excel 的指定 Sheet 页
    """
    # 写入表头
    ws.append(HEADERS)
    
    # 写入表格数据
    for table in tables:
        row_data = []
        for header in HEADERS:
            row_data.append(table.get(header, ""))  # 如果表格中没有该字段，则填充空字符串
        ws.append(row_data)

def main():
    # 获取当前目录下所有的 .docx 文件
    current_directory = os.getcwd()
    docx_files = [f for f in os.listdir(current_directory) if f.endswith(".docx") and not f.startswith("~$")]
    
    if not docx_files:
        print("当前目录下没有 .docx 文件。")
        return
    
    # 创建一个新的 Excel 工作簿
    wb = Workbook()
    
    # 删除默认创建的 Sheet 页
    default_sheet = wb.active
    wb.remove(default_sheet)
    
    # 遍历每个 .docx 文件
    for docx_file in docx_files:
        word_file_path = os.path.join(current_directory, docx_file)
        sheet_name = os.path.splitext(docx_file)[0]  # 以 .docx 文件名命名 Sheet 页
        
        # 从 Word 中提取表格数据
        tables = extract_tables_from_word(word_file_path)
        
        # 创建新的 Sheet 页
        ws = wb.create_sheet(title=sheet_name)
        
        # 将表格数据写入 Excel 的 Sheet 页
        write_tables_to_excel(tables, ws)
        
        print(f"文件 {docx_file} 的数据已成功写入 Sheet 页 {sheet_name}")
    
    # 保存 Excel 文件
    excel_file_path = os.path.join(current_directory, "output.xlsx")
    wb.save(excel_file_path)
    print(f"所有数据已成功导出到 {excel_file_path}")

if __name__ == "__main__":
    main()